#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
IR 推导式表达式节点定义
"""

from dataclasses import dataclass, field
from typing import List as TypingList, Iterator, Optional
import ast

from mcpy.ir.expressions.basic import IRName
from mcpy.ir.expressions.literals import IRTuple
from .base import IRExpr, build_expr_from_ast
from ..base import IRNode, get_default_ir_node
from ...symbol_table import Symbol, SymbolKind


@dataclass
class IRComprehension(IRExpr):
    """推导式节点"""

    target: IRNode = field(default_factory=get_default_ir_node)
    iter: IRNode = field(default_factory=get_default_ir_node)
    ifs: TypingList[IRNode] = field(default_factory=list)
    node_type: str = "comprehension"

    def creates_scope(self) -> bool:
        """推导式节点创建作用域"""
        return True

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.target
        yield self.iter
        for if_expr in self.ifs:
            yield if_expr
        yield from super().get_child_nodes()

    def get_local_symbols(self) -> Iterator[Symbol]:
        if isinstance(self.target, IRName):
            yield self.target.get_symbol()
        elif isinstance(self.target, IRTuple):
            for element in self.target.elements:
                if isinstance(element, IRName):
                    yield element.get_symbol()

    def build_from_ast(self, node: ast.comprehension) -> None:
        """从 AST 推导式节点构建 IR 推导式节点内容

        Args:
            node: AST 推导式节点
        """
        super().build_from_ast(node)
        self.target = build_expr_from_ast(node.target, self)
        self.iter = build_expr_from_ast(node.iter, self)

        # 构建条件列表
        self.ifs = []
        for if_expr in node.ifs:
            if_ir = build_expr_from_ast(if_expr, self)
            self.ifs.append(if_ir)


@dataclass
class IRGeneratorExp(IRExpr):
    """生成器表达式节点"""

    elt: IRNode = field(default_factory=get_default_ir_node)
    generators: TypingList[IRComprehension] = field(default_factory=list)
    node_type: str = "generator_exp"

    def creates_scope(self) -> bool:
        """生成器表达式创建作用域"""
        return True

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.elt
        for generator in self.generators:
            yield generator
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.GeneratorExp) -> None:
        """从 AST 生成器表达式节点构建 IR 生成器表达式节点内容

        Args:
            node: AST 生成器表达式节点
        """
        super().build_from_ast(node)
        self.elt = build_expr_from_ast(node.elt, self)

        # 构建生成器列表
        self.generators = []
        for generator in node.generators:
            generator_ir = IRComprehension(parent=self)
            generator_ir.build_from_ast(generator)
            self.generators.append(generator_ir)


@dataclass
class IRListComp(IRExpr):
    """列表推导式节点"""

    elt: IRNode = field(default_factory=get_default_ir_node)
    generators: TypingList[IRComprehension] = field(default_factory=list)
    node_type: str = "list_comp"

    def creates_scope(self) -> bool:
        """列表推导式创建作用域"""
        return True

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.elt
        for generator in self.generators:
            yield generator
        yield from super().get_child_nodes()

    def get_local_symbols(self) -> Iterator[Symbol]:
        """获取列表推导式中定义的所有本地符号

        列表推导式创建自己的作用域，包含生成器定义的目标变量

        Returns:
            本地符号的迭代器
        """
        # 收集所有生成器定义的目标变量
        for generator in self.generators:
            yield from generator.get_local_symbols()

    def build_from_ast(self, node: ast.ListComp) -> None:
        """从 AST 列表推导式节点构建 IR 列表推导式节点内容

        Args:
            node: AST 列表推导式节点
        """
        super().build_from_ast(node)
        self.elt = build_expr_from_ast(node.elt, self)

        # 构建生成器列表
        self.generators = []
        for generator in node.generators:
            generator_ir = IRComprehension(parent=self)
            generator_ir.build_from_ast(generator)
            self.generators.append(generator_ir)


@dataclass
class IRSetComp(IRExpr):
    """集合推导式节点"""

    elt: IRNode = field(default_factory=get_default_ir_node)
    generators: TypingList[IRComprehension] = field(default_factory=list)
    node_type: str = "set_comp"

    def creates_scope(self) -> bool:
        """集合推导式创建作用域"""
        return True

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.elt
        for generator in self.generators:
            yield generator
        yield from super().get_child_nodes()

    def get_local_symbols(self) -> Iterator[Symbol]:
        """获取集合推导式中定义的所有本地符号

        集合推导式创建自己的作用域，包含生成器定义的目标变量

        Returns:
            本地符号的迭代器
        """
        # 收集所有生成器定义的目标变量
        for generator in self.generators:
            yield from generator.get_local_symbols()

    def build_from_ast(self, node: ast.SetComp) -> None:
        """从 AST 集合推导式节点构建 IR 集合推导式节点内容

        Args:
            node: AST 集合推导式节点
        """
        super().build_from_ast(node)
        self.elt = build_expr_from_ast(node.elt, self)

        # 构建生成器列表
        self.generators = []
        for generator in node.generators:
            generator_ir = IRComprehension(parent=self)
            generator_ir.build_from_ast(generator)
            self.generators.append(generator_ir)


@dataclass
class IRDictComp(IRExpr):
    """字典推导式节点"""

    key: IRNode = field(default_factory=get_default_ir_node)
    value: IRNode = field(default_factory=get_default_ir_node)
    generators: TypingList[IRComprehension] = field(default_factory=list)
    node_type: str = "dict_comp"

    def creates_scope(self) -> bool:
        """字典推导式创建作用域"""
        return True

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        yield self.key
        yield self.value
        for generator in self.generators:
            yield generator
        yield from super().get_child_nodes()

    def get_local_symbols(self) -> Iterator[Symbol]:
        """获取字典推导式中定义的所有本地符号

        字典推导式创建自己的作用域，包含生成器定义的目标变量

        Returns:
            本地符号的迭代器
        """
        # 收集所有生成器定义的目标变量
        for generator in self.generators:
            yield from generator.get_local_symbols()

    def build_from_ast(self, node: ast.DictComp) -> None:
        """从 AST 字典推导式节点构建 IR 字典推导式节点内容

        Args:
            node: AST 字典推导式节点
        """
        super().build_from_ast(node)
        self.key = build_expr_from_ast(node.key, self)
        self.value = build_expr_from_ast(node.value, self)

        # 构建生成器列表
        self.generators = []
        for generator in node.generators:
            generator_ir = IRComprehension(parent=self)
            generator_ir.build_from_ast(generator)
            self.generators.append(generator_ir)
